package raft

import (
	"context"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"log"
	"time"
)

const SnapshotThreshold = 1000

func (n *RaftNode) maybeTakeSnapshot() {
	if n.lastApplied-n.lastSnapshotIndex >= SnapshotThreshold {
		n.takeSnapshot()
	}
}

func (n *RaftNode) takeSnapshot() {
	// 获取状态机快照
	snapshot, err := n.stateMachine.Snapshot()
	if err != nil {
		log.Printf("Node %d failed to create snapshot: %v", n.id, err)
		return
	}

	// 压缩日志
	n.logs = n.logs[n.lastApplied-n.lastSnapshotIndex:]
	n.lastSnapshotIndex = n.lastApplied
	n.lastSnapshotTerm = n.logTerm(n.lastApplied)

	// 保存快照
	n.persister.SaveSnapshot(snapshot, n.lastSnapshotIndex, n.lastSnapshotTerm)
	log.Printf("Node %d created snapshot at index %d", n.id, n.lastSnapshotIndex)
}

func (n *RaftNode) sendSnapshot(peerID int32) {
	snapshot, lastIndex, lastTerm, err := n.persister.LoadSnapshot()
	if err != nil || lastIndex <= 0 {
		return
	}

	args := &InstallSnapshotReq{
		Term:              n.currentTerm,
		LeaderId:          n.id,
		LastIncludedIndex: lastIndex,
		LastIncludedTerm:  lastTerm,
		Snapshot:          snapshot,
	}

	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	_, err = n.peers[peerID].client.InstallSnapshot(ctx, args)
	if err != nil {
		log.Printf("Node %d failed to send snapshot to %d: %v", n.id, peerID, err)
		return
	}

	n.mu.Lock()
	n.nextIndex[peerID] = lastIndex + 1
	n.matchIndex[peerID] = lastIndex
	n.mu.Unlock()
}

// gRPC 方法实现
func (n *RaftNode) InstallSnapshot(ctx context.Context, req *InstallSnapshotReq) (*InstallSnapshotResp, error) {
	n.mu.Lock()
	defer n.mu.Unlock()

	if req.Term < n.currentTerm {
		return &InstallSnapshotResp{Term: n.currentTerm}, nil
	}

	if req.Term > n.currentTerm {
		n.currentTerm = req.Term
		n.state = Follower
		n.votedFor = -1
	}

	n.lastHeard = time.Now()

	// 应用快照到状态机
	if err := n.stateMachine.ApplySnapshot(req.Snapshot); err != nil {
		return &InstallSnapshotResp{Term: n.currentTerm}, status.Error(codes.Internal, err.Error())
	}

	// 更新状态
	n.lastSnapshotIndex = req.LastIncludedIndex
	n.lastSnapshotTerm = req.LastIncludedTerm
	n.commitIndex = max(n.commitIndex, n.lastSnapshotIndex)
	n.lastApplied = max(n.lastApplied, n.lastSnapshotIndex)
	n.logs = []*LogEntry{{Term: n.lastSnapshotTerm, Index: n.lastSnapshotIndex}}

	// 保存状态
	n.persister.SaveState(n.currentTerm, n.votedFor, n.logs, n.lastSnapshotIndex, n.lastSnapshotTerm)
	n.persister.SaveSnapshot(req.Snapshot, n.lastSnapshotIndex, n.lastSnapshotTerm)

	return &InstallSnapshotResp{Term: n.currentTerm}, nil
}
